module efield

  use dimensions
  use fft_3d
  use fft_wrap

  real,    private, allocatable   :: kk3Dinv(:,:,:)   ! matrix of k-squared inversed
  complex, private, allocatable   :: phi0init(:,:,:)  ! initial potential of background Efield 
  complex, private, allocatable   :: phi03DFT(:,:,:)  ! FT of potential of background Efield

  complex, private, allocatable   :: phi3DFT(:,:,:)
  complex, private, allocatable   :: dummy3D(:,:,:)

  complex, private   :: dummy1Dz(jz)

  complex, private   :: E0Cinit(jz)         ! initial background E-field at the center
  complex, private   :: E0C(jz)             ! background E-field at the center
  complex, private   :: PhiInt
  complex, private   :: PhiTot

contains

!--------------------------------------------------------------------------------------
subroutine efield_init()

  use parameters, only : phi0, kz, iE0Modulation, SpeckleWidth, SpeckleDepth, iunit
  use parameters, only : iSingleMode
  use mpi_wrap

  implicit none

  integer              :: ix, iy, iz

  complex              :: xMod(jx)
  complex              :: yMod(jy)

  write(msg,*) char(10), ' Initialing efield module'
  call mpi_wrap_msg

  allocate ( kk3Dinv(jz,jy,jx) )
  allocate ( phi0init(jz,jy,jx) )
  allocate ( phi03DFT(jz,jy,jx) )
  allocate ( phi3DFT(jz,jy,jx) )
  allocate ( dummy3D(jz,jy,jx) )


  !-- setting modulation vectors --

  if (iE0Modulation) then

      xMod(:) = exp( - 20 * xvec(:)**2 / SpeckleWidth**2 )
      yMod(:) = exp( - 20 * yvec(:)**2 / SpeckleDepth**2 )

      call dealias_1D(xMod,jx)
      call dealias_1D(yMod,jy)
 
  else

     xMod(:) = 1.
     yMod(:) = 1.

  end if


  !-- initialization of background field 

  E0Cinit(:) =  -iunit * kz * phi0 *exp ( iunit * kz * zvec(:) )


  do ix=1,jx
     do iy=1,jy
        phi0init(:,iy,ix) = xMod(ix)*yMod(iy)*phi0 *exp ( iunit * kz * zvec(:)  ) 
     enddo
  enddo
  
  E0Cinit = real(E0Cinit)
  phi0init = real(phi0init)

  !-- take Fourier transform of background potential, store result

  E0C      = E0Cinit
  phiTot   = phi0*kz*kz

  phi03DFT = phi0init

  call dfftw_execute_dft(plan_xyzbackward, phi03DFT, phi03DFT)
  phi03DFT = phi03DFT/(jx*jy*jz)


  !-- ignore external potential in the linearized code --
  !
  ! (It should be added to BGK background rather than to perturbation.
  ! Also, this part needs to be reorganized if we ever go nonlinear.)

  if ((jx > 1) .or. (iSingleMode) ) phi03DFT = 0



  !-- generate matrix of k-squared

  do ix=1,jx
     do iy=1,jy
        do iz=1,jz
           kk3Dinv(iz,iy,ix) = 1./(kXvec(ix)**2 + kZvec(iz)**2  + kYvec(iy)**2)
        end do
     end do
  end do

  kk3Dinv(1,1,1)=0


end subroutine efield_init


!-------------------------------------------------------------------------------------
subroutine efield_compute_E(dfu, Ex3D, Ey3D, Ez3D, filter)
!-------------------------------------------------------------------------------------

  use parameters, only : iunit, nz, iSingleMode

  implicit none

  complex, intent(in)  ::  dfu(jv,jz,jy,jx,jdir) 
  real,    intent(in)  ::  filter(jz,jy,jx)

  complex, intent(out) ::  Ez3D(jz,jy,jx)
  complex, intent(out) ::  Ex3D(jz,jy,jx)
  complex, intent(out) ::  Ey3D(jz,jy,jx)

  !real,    intent(out) ::  OmegaBounce

  real     PhiCenter
  integer  ix, iy, iz

      
  !-- find density --

  call  sum_density(dfu, phi3DFT)


  !-- find phi(center) for diagnostics later --

  dummy1Dz = phi3DFT(:, jy/2+1, jx/2+1)
  call dfftw_execute_dft(plan_zforward, dummy1Dz, dummy1Dz)
  dummy1Dz = dummy1Dz/float(jz)
  phiInt   = 2*dummy1Dz(1+nz)
 
  !-- find FT of potential --
 
  call dfftw_execute_dft(plan_xyzforward, phi3DFT, phi3DFT)
  phi3DFT = phi3DFT/(jx*jy*jz)

  phi3DFT(:,:,:) = phi3DFT(:,:,:)*filter(:,:,:)

  phi3DFT = phi3DFT * kk3Dinv          ! inversing Laplacian in k-space

  phi3DFT = phi3DFT + phi03DFT         ! total potential


        
  !-- differentiate potential to find z-component of electric field -- 

  do iz=1,jz
     dummy3D(iz,:,:) = -iunit * kZvec(iz) * phi3DFT(iz,:,:)
  end do

  call dfftw_execute_dft(plan_xyzbackward, dummy3D, dummy3D)


  Ez3D = dummy3D



  !-- differentiate potential to find x-component of electric field -- 

  if ((jx .eq. 1) .and. .not. (iSingleMode) ) then

     Ex3D = 0.

  else

     do ix=1,jx
        dummy3D(:,:,ix) = -iunit * kXvec(ix) * phi3DFT(:,:,ix)
     end do

     call dfftw_execute_dft(plan_xyzbackward, dummy3D, dummy3D)

     Ex3D = dummy3D

  end if


  !-- differentiate potential to find y-component of electric field -- 

  if (jy .eq. 1) then

     Ey3D = 0.

  else

     do iy=1,jy
        dummy3D(:,iy,:) = -iunit * kYvec(iy) * phi3DFT(:,iy,:)
     end do

     call dfftw_execute_dft(plan_xyzbackward, dummy3D, dummy3D)

     Ey3D = dummy3D

  end if



  !-- find phi(center) for diagnostics later --

  !dummy1Dz(:) = Ez3D(:, jy/2+1, jx/2+1)
  !call dfftw_execute_dft(plan_zforward, dummy1Dz, dummy1Dz)
  !dummy1Dz = dummy1Dz/float(jz)

  ! amplitude of TOTAL potential at speckle center of sin(kz*z + phase)
  !PhiCenter = 2 * abs(dummy1Dz(nz+1) / kZvec(nz+1))

  !OmegaBounce = kZvec(nz+1) * sqrt(PhiCenter)
     
  return

end subroutine efield_compute_E
    

!-------------------------------------------------------------------------------------
subroutine efield_evolve_E0(t)                                   ! input
!-------------------------------------------------------------------------------------

  use parameters, only : kz, phi0, vPhase, iWaveFrame, iunit

  implicit none

  real, intent(in)  :: t    
  complex           :: ExpKVt

  if ( iWaveFrame==1 ) return


  ExpKVt = exp( - iunit*t*vPhase*kz  )
     
  E0C      = ExpKVt * E0Cinit         
  phi03DFT = ExpKVt * phi0init
  phiTot   = ExpKVt * phi0*kz*kz
     
  !-- take Fourier transform of initial potential, store result

  call dfftw_execute_dft(plan_xyzbackward, phi03DFT, phi03DFT)


end subroutine efield_evolve_E0


!-------------------------------------------------------------------------------------
subroutine efield_get_epsilon(epsilon)                ! output 
!-------------------------------------------------------------------------------------

  implicit none

  complex, intent(out)    :: epsilon

  epsilon = phiTot / (phiInt + phiTot)
 
end subroutine efield_get_epsilon

!-------------------------------------------------------------------------------------
subroutine efield_get_E0z_center(E0zCenter)            ! output
!-------------------------------------------------------------------------------------

  implicit none

  real, intent(out)       :: E0zCenter(jz)

  E0zCenter = real(E0C)

end subroutine efield_get_E0z_center



!-------------------------------------------------------------------------------------
subroutine sum_density(dfu, dens)
!-------------------------------------------------------------------------------------

  use mpi_wrap

  implicit none

  complex, intent(in)    :: dfu(jv, jz, jy, jx, jDir)
  complex, intent(out)   :: dens(jz, jy, jx)

  integer                :: ix, iy, iz

  do ix=1,jx
     do iy=1,jy
        do iz=1,jz
           dummy3D(iz,iy,ix) = sum(dfu(:,iz,iy,ix,:))*dv 
        end do
     end do
  end do

  call mpi_allreduce(dummy3D, dens, 2*jz*jy*jx, &
      MPI_REAL8, MPI_SUM, MPI_COMM_WORLD, mpi_err)


end subroutine sum_density


!-------------------------------------------------------------------------------------

  
end module
